import numpy as np
import pandas as pd
from sklearn import neighbors
from sklearn.model_selection import train_test_split


class KNN:
    def __init__(self, dataset, args):
        self.input_dim = args.get('input_dim', 9)
        self.output_dim = args.get('output_dim', 1)
        self.n_neighbors = args.get('n_neighbors', 10)
        self.h = args.get('h', .02)
        self.weights = args.get('weights', 'uniform')

        df = pd.read_csv(dataset, delimiter=',', header=None)
        self.X, self.y = df.values[:, :self.input_dim], df.values[:, -self.output_dim]
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(self.X, self.y, random_state=666)

    def run(self):
        clf = neighbors.KNeighborsClassifier(self.n_neighbors, weights=self.weights)
        clf.fit(self.X_train, self.y_train)

        predict = clf.predict(self.X_test)
        all_data = np.hstack((self.X_test, self.y_test.reshape(-1, 1), predict.reshape(-1, 1)))

        cols = ["fat", "protein", "snf", "acidity", "lead", "mercury", "arsenic", "chromium", "aflatoxins", "truth", "predict"]
        raw_data = [{k: v for k, v in zip(cols, row)} for row in all_data.tolist()]

        return {
            'TrainAccuracy': clf.score(self.X_train, self.y_train),
            'TestAccuracy': clf.score(self.X_test, self.y_test),
            'Truth': list(map(int, list(self.y_test))),
            'Predict': list(map(int, list(predict))),
            'RawData': raw_data,
            'Description': 'KNN算法的核心思想是，如果一个样本在特征空间中的K个最相邻的样本中的大多数属于某一个类别，则该样本也属于这个类别，并具有这个类别上样本的特性。'
                           '该方法在确定分类决策上只依据最邻近的一个或者几个样本的类别来决定待分样本所属的类别。KNN方法在类别决策时，只与极少量的相邻样本有关。'
                           '由于KNN方法主要靠周围有限的邻近的样本，而不是靠判别类域的方法来确定所属类别的，因此对于类域的交叉或重叠较多的待分样本集来说，KNN方法较其他方法更为适合。'
                           '例如，探讨食品风险因素，并根据风险因素预测食品发生风险的概率或等级。自变量既可以是连续的，也可以是分类的。然后通过KNN，根据相邻样本对食品风险等级进行分类。'
        }

        # # 模型测试
        # with open('../../result/knn_%s.json' % self.dataset, 'w', encoding='utf8') as f:
        #     json.dump({'Train accuracy': clf.score(self.X_train, self.y_train),
        #                'Test accuracy': clf.score(self.X_test, self.y_test),
        #                'Truth': list(map(int, list(self.y))),
        #                'Predict': list(map(int, list(clf.predict(self.X))))},
        #               f)

        # np.savetxt('../../result/knn_%s.txt' % self.dataset, np.c_[self.X, clf.predict(self.X)], fmt='%.06f')

        # values = clf.predict(self.X)
        #
        # plt.rcParams['font.family'] = ['sans-serif']
        # plt.rcParams['font.sans-serif'] = ['SimHei']
        # levels = ['IV等品', 'V等品', 'III等品', 'II等品', 'I等品']
        # b_1 = [Counter(self.y)[i] for i in range(5)]
        # b_2 = [Counter(values)[i] for i in range(5)]
        #
        # bar_width = 0.2
        # bar_1 = list(range(len(levels)))
        # bar_2 = [i+bar_width for i in bar_1]
        #
        # # 导入数据，绘制条形图
        # plt.bar(bar_1, b_1, width=bar_width, label='真实')
        # plt.bar(bar_2, b_2, width=bar_width, label='预测')
        #
        # # 添加标题
        # plt.title('食品质量分类预测', size=20)
        # # 添加xy轴
        # plt.xlabel('食品等级')
        # plt.ylabel('数量')
        # # x轴刻度
        # plt.xticks(range(len(levels)), levels)
        # plt.legend()
        #
        # # 展示效果图
        # # plt.show()
        # plt.savefig('../../result/knn_%s.png' % self.dataset)


if __name__ == '__main__':
    print(KNN('../../data/sterilizedmilk_for_classfication.csv', {}).run())
